#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import json
import re
import argparse
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sqlalchemy import create_engine
from collections import Counter
from wordcloud import WordCloud
import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from sklearn.feature_extraction.text import TfidfVectorizer
from matplotlib.font_manager import FontProperties

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置显示中文字体
plt.rcParams['axes.unicode_minus'] = False    # 设置正常显示符号

# 下载必要的NLTK资源
def download_nltk_resources():
    """下载NLTK所需资源"""
    resources = ['punkt', 'stopwords', 'wordnet']
    for resource in resources:
        try:
            nltk.download(resource, quiet=True)
        except:
            print(f"无法下载NLTK资源: {resource}")

def load_data(data_dir):
    """从SQLite数据库加载数据"""
    db_file = os.path.join(data_dir, "prompts.db")
    
    if not os.path.exists(db_file):
        raise FileNotFoundError(f"数据库文件 {db_file} 不存在，请先运行collect_prompt_details.py")
    
    engine = create_engine(f"sqlite:///{db_file}")
    df = pd.read_sql("SELECT * FROM prompts", engine)
    
    # 加载卖家信息
    seller_info_file = os.path.join(data_dir, "seller_info.json")
    with open(seller_info_file, 'r', encoding='utf-8') as f:
        seller_info = json.load(f)
    
    return df, seller_info

def clean_text(text):
    """清洗文本数据"""
    if not isinstance(text, str):
        return ""
    
    # 转为小写
    text = text.lower()
    
    # 移除URLs
    text = re.sub(r'http\S+', '', text)
    
    # 移除特殊字符
    text = re.sub(r'[^\w\s]', ' ', text)
    
    # 移除多余空格
    text = re.sub(r'\s+', ' ', text).strip()
    
    return text

def preprocess_text(text, remove_stopwords=True):
    """文本预处理：分词、去停用词、词形还原"""
    if not isinstance(text, str) or not text:
        return []
    
    # 清洗文本
    text = clean_text(text)
    
    # 分词
    tokens = word_tokenize(text)
    
    # 去停用词
    if remove_stopwords:
        stop_words = set(stopwords.words('english'))
        tokens = [token for token in tokens if token not in stop_words and len(token) > 2]
    
    # 词形还原
    lemmatizer = WordNetLemmatizer()
    tokens = [lemmatizer.lemmatize(token) for token in tokens]
    
    return tokens

def analyze_prompt_attributes(df, output_dir):
    """分析prompt属性分布"""
    print("分析prompt属性分布...")
    
    # 创建结果目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 设置图表样式
    sns.set(style="whitegrid")
    
    # 1. 生成价格分布直方图
    plt.figure(figsize=(10, 6))
    sns.histplot(df['price'].dropna(), bins=20, kde=True)
    plt.title("Prompt价格分布")
    plt.xlabel("价格 ($)")
    plt.ylabel("数量")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "price_distribution.png"), dpi=300)
    plt.close()
    
    # 2. 生成评分分布饼图
    plt.figure(figsize=(10, 6))
    rating_counts = df['rating'].dropna().map(lambda x: round(x)).value_counts().sort_index()
    plt.pie(rating_counts, labels=rating_counts.index, autopct='%1.1f%%', startangle=90)
    plt.title("Prompt评分分布")
    plt.axis('equal')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "rating_distribution.png"), dpi=300)
    plt.close()
    
    # 3. 分析标签分布
    all_tags = []
    for tags_str in df['tags'].dropna():
        if isinstance(tags_str, str):
            all_tags.extend([tag.strip() for tag in tags_str.split(',')])
    
    tag_counts = Counter(all_tags).most_common(15)
    tags_df = pd.DataFrame(tag_counts, columns=['Tag', 'Count'])
    
    plt.figure(figsize=(12, 6))
    sns.barplot(x='Count', y='Tag', data=tags_df)
    plt.title("最常见的15个标签")
    plt.xlabel("数量")
    plt.ylabel("标签")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "tag_distribution.png"), dpi=300)
    plt.close()
    
    # 4. 分析模型分布
    model_counts = df['model'].value_counts().head(10)
    plt.figure(figsize=(10, 6))
    model_counts.plot(kind='bar')
    plt.title("最常用的10个模型")
    plt.xlabel("模型")
    plt.ylabel("数量")
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "model_distribution.png"), dpi=300)
    plt.close()
    
    # 5. 分析类别分布
    category_counts = df['category'].value_counts().head(10)
    plt.figure(figsize=(10, 6))
    category_counts.plot(kind='pie', autopct='%1.1f%%')
    plt.title("Prompt类别分布")
    plt.axis('equal')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "category_distribution.png"), dpi=300)
    plt.close()
    
    # 6. 分析字数分布
    plt.figure(figsize=(10, 6))
    sns.histplot(df['word_count'].dropna(), bins=20, kde=True)
    plt.title("Prompt字数分布")
    plt.xlabel("字数")
    plt.ylabel("数量")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "word_count_distribution.png"), dpi=300)
    plt.close()
    
    print("prompt属性分布分析完成，图表已保存到 " + output_dir)

def analyze_content_style(df, seller_info, output_dir):
    """分析卖家内容风格"""
    print("分析卖家内容风格...")
    
    # 预处理所有标题和描述
    titles = df['name'].dropna().apply(clean_text).tolist()
    descriptions = df['description'].dropna().apply(clean_text).tolist()
    
    # 合并所有文本
    all_text = ' '.join(titles + descriptions)
    
    # 生成词云
    wordcloud = WordCloud(
        width=800, 
        height=400, 
        background_color='white',
        max_words=100,
        collocations=False
    ).generate(all_text)
    
    plt.figure(figsize=(12, 6))
    plt.imshow(wordcloud, interpolation='bilinear')
    plt.axis("off")
    plt.title("Prompt标题和描述的词云")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "content_wordcloud.png"), dpi=300)
    plt.close()
    
    # 使用TF-IDF提取关键词
    vectorizer = TfidfVectorizer(
        max_features=20,
        stop_words='english',
        ngram_range=(1, 2)
    )
    
    try:
        tfidf_matrix = vectorizer.fit_transform(descriptions)
        feature_names = vectorizer.get_feature_names_out()
        
        # 计算每个特征的平均TF-IDF得分
        mean_tfidf = np.array(tfidf_matrix.mean(axis=0)).flatten()
        
        # 创建特征和得分的DataFrame
        tfidf_df = pd.DataFrame({'feature': feature_names, 'tfidf': mean_tfidf})
        tfidf_df = tfidf_df.sort_values(by='tfidf', ascending=False).head(15)
        
        plt.figure(figsize=(12, 6))
        sns.barplot(x='tfidf', y='feature', data=tfidf_df)
        plt.title("描述中最重要的15个词/短语 (TF-IDF)")
        plt.xlabel("平均TF-IDF得分")
        plt.ylabel("词/短语")
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, "tfidf_keywords.png"), dpi=300)
        plt.close()
    except:
        print("TF-IDF分析失败，可能是因为描述文本太少")
    
    # 分析卖家Bio
    if 'bio' in seller_info and seller_info['bio']:
        bio_text = clean_text(seller_info['bio'])
        
        # 生成卖家Bio词云
        if len(bio_text) > 20:  # 确保有足够的文本生成词云
            wordcloud = WordCloud(
                width=800, 
                height=400, 
                background_color='white',
                max_words=50,
                collocations=False
            ).generate(bio_text)
            
            plt.figure(figsize=(12, 6))
            plt.imshow(wordcloud, interpolation='bilinear')
            plt.axis("off")
            plt.title("卖家Bio词云")
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, "seller_bio_wordcloud.png"), dpi=300)
            plt.close()
    
    print("卖家内容风格分析完成，图表已保存到 " + output_dir)

def analyze_engagement_metrics(df, output_dir):
    """分析参与度指标"""
    print("分析参与度指标...")
    
    # 1. 分析点赞数分布
    plt.figure(figsize=(10, 6))
    sns.histplot(df['likes'].dropna(), bins=20, kde=True)
    plt.title("Prompt点赞数分布")
    plt.xlabel("点赞数")
    plt.ylabel("数量")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "likes_distribution.png"), dpi=300)
    plt.close()
    
    # 2. 分析评分和点赞的关系
    plt.figure(figsize=(10, 6))
    sns.scatterplot(x='rating', y='likes', data=df)
    plt.title("评分与点赞数的关系")
    plt.xlabel("评分")
    plt.ylabel("点赞数")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "rating_vs_likes.png"), dpi=300)
    plt.close()
    
    # 3. 分析价格和参与度的关系
    plt.figure(figsize=(10, 6))
    sns.scatterplot(x='price', y='likes', data=df)
    plt.title("价格与点赞数的关系")
    plt.xlabel("价格 ($)")
    plt.ylabel("点赞数")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "price_vs_likes.png"), dpi=300)
    plt.close()
    
    # 4. 分析不同类别的平均参与度
    category_engagement = df.groupby('category')[['likes', 'rating']].mean().reset_index()
    category_engagement = category_engagement.sort_values(by='likes', ascending=False).head(10)
    
    plt.figure(figsize=(12, 6))
    sns.barplot(x='likes', y='category', data=category_engagement)
    plt.title("不同类别的平均点赞数")
    plt.xlabel("平均点赞数")
    plt.ylabel("类别")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "category_avg_likes.png"), dpi=300)
    plt.close()
    
    # 5. 分析不同模型的平均参与度
    model_engagement = df.groupby('model')[['likes', 'rating']].mean().reset_index()
    model_engagement = model_engagement.sort_values(by='likes', ascending=False).head(10)
    
    plt.figure(figsize=(12, 6))
    sns.barplot(x='likes', y='model', data=model_engagement)
    plt.title("不同模型的平均点赞数")
    plt.xlabel("平均点赞数")
    plt.ylabel("模型")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "model_avg_likes.png"), dpi=300)
    plt.close()
    
    print("参与度指标分析完成，图表已保存到 " + output_dir)

def generate_analysis_report(df, seller_info, output_dir):
    """生成数据分析报告"""
    report_file = os.path.join(output_dir, "analysis_report.md")
    
    with open(report_file, 'w', encoding='utf-8') as f:
        # 报告标题
        f.write("# Prompt Marketplace 卖家分析报告\n\n")
        
        # 卖家基本信息
        f.write("## 卖家基本信息\n\n")
        f.write(f"- **用户名**: {seller_info.get('username', 'N/A')}\n")
        if 'bio' in seller_info and seller_info['bio']:
            f.write(f"- **个人简介**: {seller_info['bio']}\n")
        
        # 写入统计数据
        if 'stats' in seller_info:
            f.write("- **统计数据**:\n")
            for key, value in seller_info['stats'].items():
                f.write(f"  - {key}: {value}\n")
        
        # 其他信息
        if 'other_info' in seller_info:
            f.write("- **其他信息**:\n")
            for key, value in seller_info['other_info'].items():
                f.write(f"  - {key}: {value}\n")
        
        f.write("\n")
        
        # 数据概览
        f.write("## 数据概览\n\n")
        f.write(f"- **收集的prompt总数**: {len(df)}\n")
        f.write(f"- **平均价格**: ${df['price'].mean():.2f}\n")
        f.write(f"- **平均评分**: {df['rating'].mean():.2f}\n")
        f.write(f"- **平均点赞数**: {df['likes'].mean():.2f}\n")
        f.write(f"- **平均字数**: {df['word_count'].mean():.2f}\n")
        f.write("\n")
        
        # 主要分析结果
        f.write("## 主要分析结果\n\n")
        
        # 1. Prompt属性分布
        f.write("### Prompt属性分布\n\n")
        f.write("- **价格分布**: 见图表 `price_distribution.png`\n")
        f.write("- **评分分布**: 见图表 `rating_distribution.png`\n")
        f.write("- **标签分布**: 见图表 `tag_distribution.png`\n")
        f.write("- **模型分布**: 见图表 `model_distribution.png`\n")
        f.write("- **类别分布**: 见图表 `category_distribution.png`\n")
        f.write("- **字数分布**: 见图表 `word_count_distribution.png`\n")
        f.write("\n")
        
        # 2. 内容风格分析
        f.write("### 内容风格分析\n\n")
        f.write("- **关键词**: 见图表 `content_wordcloud.png` 和 `tfidf_keywords.png`\n")
        if 'bio' in seller_info and seller_info['bio']:
            f.write("- **卖家Bio分析**: 见图表 `seller_bio_wordcloud.png`\n")
        f.write("\n")
        
        # 3. 参与度指标分析
        f.write("### 参与度指标分析\n\n")
        f.write("- **点赞分布**: 见图表 `likes_distribution.png`\n")
        f.write("- **评分与点赞关系**: 见图表 `rating_vs_likes.png`\n")
        f.write("- **价格与点赞关系**: 见图表 `price_vs_likes.png`\n")
        f.write("- **类别参与度**: 见图表 `category_avg_likes.png`\n")
        f.write("- **模型参与度**: 见图表 `model_avg_likes.png`\n")
        f.write("\n")
        
        # 4. 结论和见解
        f.write("## 结论和见解\n\n")
        f.write("基于数据分析，以下是对此卖家的主要观察结果：\n\n")
        
        # 这里根据实际数据添加一些见解，例如：
        
        # 价格策略
        avg_price = df['price'].mean()
        market_position = "中等" if 5 <= avg_price <= 15 else ("高端" if avg_price > 15 else "低端")
        f.write(f"1. **价格策略**: 该卖家的prompt平均价格为${avg_price:.2f}，处于{market_position}市场。\n")
        
        # 内容类型偏好
        top_categories = df['category'].value_counts().head(3).index.tolist()
        if top_categories:
            f.write(f"2. **内容偏好**: 该卖家主要专注于以下类别的prompt: {', '.join(top_categories)}。\n")
        
        # 受欢迎度分析
        avg_likes = df['likes'].mean()
        popularity = "较高" if avg_likes > 10 else ("中等" if 5 <= avg_likes <= 10 else "较低")
        f.write(f"3. **受欢迎度**: 该卖家的prompt平均获得{avg_likes:.2f}个点赞，受欢迎度{popularity}。\n")
        
        # 质量评估
        avg_rating = df['rating'].mean()
        quality = "优秀" if avg_rating > 4.5 else ("良好" if 4.0 <= avg_rating <= 4.5 else ("一般" if 3.5 <= avg_rating < 4.0 else "较差"))
        if not pd.isna(avg_rating):
            f.write(f"4. **质量评估**: 基于{avg_rating:.2f}的平均评分，该卖家的prompt质量整体{quality}。\n")
        
        f.write("\n")
        
        # 结束语
        f.write("## 附录：数据可视化图表\n\n")
        f.write("所有生成的数据可视化图表已保存在同一目录下。\n")
    
    print(f"分析报告已生成: {report_file}")

def main():
    parser = argparse.ArgumentParser(description="分析Prompt Marketplace数据")
    parser.add_argument("--data_dir", default="../data", help="数据目录路径")
    parser.add_argument("--output_dir", default="../results", help="结果输出目录路径")
    args = parser.parse_args()
    
    try:
        # 下载NLTK资源
        download_nltk_resources()
        
        # 加载数据
        df, seller_info = load_data(args.data_dir)
        print(f"加载了 {len(df)} 条prompt记录")
        
        # 分析prompt属性
        analyze_prompt_attributes(df, args.output_dir)
        
        # 分析内容风格
        analyze_content_style(df, seller_info, args.output_dir)
        
        # 分析参与度指标
        analyze_engagement_metrics(df, args.output_dir)
        
        # 生成分析报告
        generate_analysis_report(df, seller_info, args.output_dir)
        
        print(f"数据分析完成！结果已保存到 {args.output_dir}")
    
    except Exception as e:
        print(f"数据分析时发生错误: {str(e)}")

if __name__ == "__main__":
    main() 